package com.dynamic.mybatis.core.datasource;


import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.io.Closeable;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

public class DynamicMultipleDataSource extends AbstractRoutingDataSource {

    private Map<Object,Object> dataSources;

    public DynamicMultipleDataSource(DataSource defaultDataSource, Map<Object, Object> dataSources) {
        super.setDefaultTargetDataSource(defaultDataSource);
        super.setTargetDataSources(dataSources);
        super.afterPropertiesSet();
    }

    public void set(Map<Object, Object> dataSources) {
        if (dataSources == null || dataSources.size() < 1) {
            return;
        }
        close(this.dataSources);
        this.dataSources = dataSources;
        super.setTargetDataSources(dataSources);
        super.afterPropertiesSet();
    }

    public void put(Object key,DataSource dataSource) {
        if (dataSources == null || dataSources.size() < 1) {
             this.dataSources = new HashMap<>();
        }
        this.dataSources.put(key,dataSource);
        set(this.dataSources);
    }

    public void remove(Object key) {
        if (dataSources == null || dataSources.size() < 1) {
           return;
        }
        close(key);
        this.dataSources.remove(key);
        set(this.dataSources);
    }

    private void close(Map<Object,Object> dataSources){
        //被重置数据源和重置数据源属于同一对象时，不进行关闭
        if(dataSources == null || dataSources.size() < 1  || (this.dataSources != null && this.dataSources == dataSources)){
            return;
        }
        dataSources.forEach((key,value) ->{
            DataSource dataSource1 = this.resolveSpecifiedDataSource(dataSources.get(key));
            DataSource dataSource2 = this.resolveSpecifiedDataSource(this.dataSources.get(key));
            if(dataSource1 != null && dataSource2 != null && dataSource1 != dataSource2){
                close(key);
            }
        });
    }

    private void close(Object key){
        DataSource dataSource = this.resolveSpecifiedDataSource(this.dataSources.get(key));
        if(dataSource != null && dataSource instanceof Closeable){
            try {
                ((Closeable)dataSource).close();
            } catch (IOException e) {
                throw new RuntimeException("dataSource close is error."+e,e);
            }
        }
    }

    @Override
    protected Object determineCurrentLookupKey() {
        return DynamicDataSourceHolder.get();
    }


}
